package edu.zju.gis.dbfg.model.util;

import org.gdal.gdal.gdal;
import org.gdal.ogr.*;
import org.gdal.osr.SpatialReference;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Objects;

/**
 * <p>transform a shapefile, FileGDB, Personal GDB, Geomedia Database,
 * or MS Access DB to wkt string, and write it to text file.</p>
 *
 * @author yanlo yanlong_lee@qq.com
 * @version 1.1 2017/10/18
 */
public class WKTWriter {

    private static final String DEFAULT_EXTENSION = ".wkt";

    private String file;
    private SpatialReference targetSpatialReference;
    private String extension = DEFAULT_EXTENSION;
    private boolean outMetadata;

    static {
        gdal.AllRegister();
        ogr.RegisterAll();
        gdal.SetConfigOption("GDAL_FILENAME_IS_UTF8", "YES");
        gdal.SetConfigOption("SHAPE_ENCODING", "GB2312");
    }

    /**
     * @param file path of spatial file(e.g. shapefile, gdb)
     */
    public WKTWriter(String file) {
        this(file, null, false);
    }

    public WKTWriter(String file, boolean outMetadata) {
        this(file, null, outMetadata);
    }

    public WKTWriter(String file, SpatialReference targetSptialReference) {
        this(file, targetSptialReference, false);
    }

    public WKTWriter(String file, SpatialReference targetSptialReference, boolean outMetadata) {
        this.file = file;
        this.targetSpatialReference = targetSptialReference;
        this.outMetadata = outMetadata;
    }

    public String getExtension() {
        return this.extension;
    }

    public void setExtension(String extension) {
        this.extension = extension;
    }

    /**
     * write the input file to the <b>outputDir</b>. If the output file already exists, it will be deleted.
     *
     * @param outputDir the output directory
     * @see #write(File)
     */
    public void write(String outputDir) throws IOException {
        File dir = new File(outputDir);
        write(dir);
    }

    /**
     * write the input file to the <b>outputDir</b>. If the output file already exists, it will be deleted.
     *
     * @param outputDir the output directory
     */
    public void write(File outputDir) throws IOException {
        if (outputDir.exists())
            deletePath(outputDir);
        outputDir.mkdirs();
        DataSource ds = ogr.Open(file);
        ds = Objects.requireNonNull(ds);
        try {
            int cnt = ds.GetLayerCount();
            for (int i = 0; i < cnt; i++) {
                Layer lyr = ds.GetLayer(i);
                String[] fields = null;
                String srWkt = null;
                if (outMetadata) {
                    fields = getFields(lyr);
                    srWkt = getSpatialReferenceAsWkt(lyr);
                }
                List<String> wkt = stringify(lyr, targetSpatialReference);
                writeTo(lyr.GetName(), fields, srWkt, wkt, outputDir);
            }
        } catch (Exception ex) {
            ex.printStackTrace();
            throw ex;
        } finally {
            ds.delete();
        }
    }

    public static void deletePath(File path) throws IOException {
        if (path.exists()) {
            if (path.isFile())
                Files.delete(path.toPath());
            else {
                File[] subPaths = path.listFiles();
                if (subPaths != null)
                    for (File subPath : subPaths)
                        deletePath(subPath);
                Files.delete(path.toPath());
            }
        }
    }

    /**
     * <p>
     * write the layer in the wkt form to the directory. A wkt file with a name of the layer's name and a corresponding field description file will be created.
     * </p>
     *
     * @param lyrName   as the output file's name
     * @param directory the directory
     */
    private void writeTo(String lyrName, List<String> wkt, File directory) throws IOException {
        writeTo(lyrName, null, null, wkt, directory);
    }

    private void writeTo(String lyrName, String[] fields, String srWkt, List<String> wkt, File directory) throws IOException {
        if (fields != null) {
            try {
                Path metaPath = Paths.get(directory.getAbsolutePath(), lyrName + "_meta.txt");
                Files.write(metaPath, Arrays.asList(String.join("\t", fields), srWkt));
            } catch (IOException e) {
                e.printStackTrace();
                throw e;
            }
        }

        File outputFile = new File(directory, lyrName + extension);
        if (outputFile.exists() && outputFile.isFile())
            outputFile.delete();
        try (FileWriter fw = new FileWriter(outputFile)) {
            for (String line : wkt) {
                fw.append(line).append("\n").flush();
            }
        } catch (IOException e) {
            e.printStackTrace();
            throw e;
        }
    }

    /**
     * @param layer input layer
     * @return a Tuple2 object that contains an array of field names(fid included) and an array of the features' wkt.
     * @see #stringify(Layer, SpatialReference)
     */
    public static List<String> stringify(Layer layer) {
        return stringify(layer, null);
    }

    /**
     * @param layer    input layer
     * @param targetSR the target SpatialReference for wkt, use the layer's sr when it is null.
     * @return a Tuple2 object that contains an array of field names(fid included) and an array of the features' wkt.
     */
    public static List<String> stringify(Layer layer, SpatialReference targetSR) {
        FeatureDefn ftd = layer.GetLayerDefn();
        int cnt = ftd.GetFieldCount();
        Feature ft;
        Geometry g;
        layer.ResetReading();
        ArrayList<String> features = new ArrayList<>((int) layer.GetFeatureCount());
        StringBuilder sb = new StringBuilder(1024);
        while ((ft = layer.GetNextFeature()) != null) {
            sb.append(ft.GetFID()).append("\t");
            for (int i = 0; i < cnt; i++)
                sb.append(ft.GetFieldAsString(i).replaceAll("[\r\n]", "").replace("\t", " ")).append("\t");
            g = ft.GetGeometryRef();
            if (g != null) {
                if (targetSR != null)
                    g.TransformTo(targetSR);
                sb.append(g.ExportToWkt());
            }
            features.add(sb.toString());
            sb.delete(0, sb.length());
        }
        return features;
    }

    /**
     * get the field name array from a layer with FID as the first field and SHAPE as the last.
     *
     * @return an array of field names
     */
    public static String[] getFields(Layer layer) {
        return getFields(layer, true, true);
    }

    /**
     * get the field name array from a layer with FID as the first field and SHAPE as the last.
     *
     * @return an array of field names
     */
    public static String[] getFields(Layer layer, boolean withFid, boolean withShape) {
        FeatureDefn ftd = layer.GetLayerDefn();
        int fieldCount = ftd.GetFieldCount();
        List<String> fields = new ArrayList<>();
        if (withFid)
            fields.add("FID");
        for (int i = 0; i < fieldCount; i++)
            if (ftd.GetFieldDefn(i).GetName().toUpperCase().equals("\t"))
                fields.add(null);
            else
                fields.add(ftd.GetFieldDefn(i).GetName().toUpperCase());
        if (withShape)
            fields.add("SHAPE");
        return fields.toArray(new String[0]);
    }

    public static String getSpatialReferenceAsWkt(Layer layer) {
        SpatialReference sr = layer.GetSpatialRef();
        return sr == null ? null : sr.ExportToWkt().replaceAll("[\r\n]", "");
    }

}
